import math
from typing import Dict, Tuple, List
from copy import copy


class VehiclePerceptionPredictor:
    def __init__(self, min_duration=2.0, frame_interval=0.1, predict_frames=20):
        self.min_duration = min_duration
        self.frame_interval = frame_interval
        self.predict_frames = predict_frames
        self.METERS_PER_DEGREE_LAT = 111111
        self.METERS_PER_DEGREE_LON = 111320

    def is_valid_target(self, timestamps: List[int]) -> bool:
        if not timestamps:
            return False
        duration = (timestamps[-1] - timestamps[0]) / 1000.0
        return duration >= self.min_duration

    def calculate_speed_and_heading(self, positions: List[Tuple[float, float, int]], n_frames=5) -> Tuple[float, float]:
        recent_positions = positions[-n_frames:]
        if len(recent_positions) < 2:
            return 0.0, 0.0

        total_distance = 0.0
        total_time = 0.0
        prev_lon, prev_lat, prev_ts = None, None, None

        for lon, lat, ts in recent_positions:
            if prev_lon is not None:
                dx = (lon - prev_lon) * self.METERS_PER_DEGREE_LON * math.cos(math.radians(lat))
                dy = (lat - prev_lat) * self.METERS_PER_DEGREE_LAT
                distance = math.sqrt(dx ** 2 + dy ** 2)
                time_diff = (ts - prev_ts) / 1000.0
                total_distance += distance
                total_time += time_diff
            prev_lon, prev_lat, prev_ts = lon, lat, ts

        if total_time == 0:
            return 0.0, 0.0

        avg_speed = total_distance / total_time

        if len(recent_positions) >= 2:
            lon1, lat1, _ = recent_positions[-2]
            lon2, lat2, _ = recent_positions[-1]
            dx = (lon2 - lon1) * self.METERS_PER_DEGREE_LON * math.cos(math.radians(lat2))
            dy = (lat2 - lat1) * self.METERS_PER_DEGREE_LAT
            heading = math.atan2(dy, dx)
        else:
            heading = 0.0

        return avg_speed, heading

    def predict_future_positions(self, last_data: Dict, last_ts: int, speed: float, heading: float) -> Dict[int, Dict]:
        predictions = {}
        last_lon = last_data['lon']
        last_lat = last_data['lat']
        cos_heading = math.cos(heading)
        sin_heading = math.sin(heading)
        cos_lat = math.cos(math.radians(last_lat))
        lon_scale = self.METERS_PER_DEGREE_LON * cos_lat
        lat_scale = self.METERS_PER_DEGREE_LAT

        for i in range(1, self.predict_frames + 1):
            pred_ts = last_ts + int(i * self.frame_interval * 1000)
            pred_data = copy(last_data)
            distance = speed * self.frame_interval * i
            delta_lon = (distance * sin_heading) / lon_scale
            delta_lat = (distance * cos_heading) / lat_scale
            pred_data['lon'] = last_lon + delta_lon
            pred_data['lat'] = last_lat + delta_lat
            pred_data['is_predicted'] = True
            pred_data['predicted_frame'] = i
            pred_data['predicted_time_offset'] = i * self.frame_interval
            predictions[pred_ts] = pred_data

        return predictions

    def process_vehicle_data(self, per_car_info_dict: Dict[str, Dict[str, Dict[int, Dict]]]) -> Dict[
        str, Dict[str, Dict[int, Dict]]]:
        predictions = {}
        valid_targets = {}

        # 关键修复：遍历前将字典转为列表，避免迭代时字典被修改
        for sourceId in list(per_car_info_dict.keys()):
            car_dict = per_car_info_dict[sourceId]
            valid_targets[sourceId] = {}
            for car_id in list(car_dict.keys()):
                ts_data = car_dict[car_id]
                # 复制ts_data的副本进行迭代，避免原字典被修改
                ts_data_copy = dict(ts_data.items())  # 或使用ts_data.copy()
                timestamps = sorted(ts_data_copy.keys())
                if self.is_valid_target(timestamps):
                    # 使用副本生成positions，避免直接操作原字典
                    positions = [(data['lon'], data['lat'], ts) for ts, data in ts_data_copy.items()]
                    positions.sort(key=lambda x: x[2])
                    valid_targets[sourceId][car_id] = {
                        'positions': positions,
                        'timestamps': timestamps,
                        'last_data': ts_data_copy[timestamps[-1]],  # 使用副本获取最后一条数据
                        'last_ts': timestamps[-1]
                    }

        for sourceId, car_dict in valid_targets.items():
            source_predictions = {}
            for car_id, target_data in car_dict.items():
                speed, heading = self.calculate_speed_and_heading(target_data['positions'])
                future_positions = self.predict_future_positions(
                    target_data['last_data'],
                    target_data['last_ts'],
                    speed,
                    heading
                )
                source_predictions[car_id] = future_positions
            if source_predictions:
                predictions[sourceId] = source_predictions

        return predictions